from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
import torch
from diffusers import FluxPipeline
import os
import random
import datetime
import configparser  # 导入 configparser

app = FastAPI()

# 读取配置文件
config = configparser.ConfigParser()
config.read('config.ini')
base_url = config['DEFAULT']['BASE_URL']
model_path = config['DEFAULT']['MODEL_PATH']
torch_dtype = getattr(torch, config['DEFAULT']['TORCH_DTYPE'])

# 加载模型到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = FluxPipeline.from_pretrained(model_path, torch_dtype=torch_dtype).to(device)

# 定义请求体的数据模型
class ImageRequest(BaseModel):
    prompt: str
    image_size: str
    num_inference_steps: int

@app.post("/FLUX.1-schnell/text-to-image")
async def generate_image(request: ImageRequest):
    # 验证 image_size 格式，支持 'x'、'*'、'X' 和 ',' 分隔符
    try:
        # 替换 '*' 和 ',' 为 'x'，并将 'X' 转为小写
        image_size_corrected = request.image_size.replace('*', 'x').replace(',', 'x').replace('X', 'x')
        width, height = map(int, image_size_corrected.split('x'))
        
        # 调整 width 和 height 为最接近的 8 的倍数
        width = (width // 8) * 8
        height = (height // 8) * 8
        
        # 确保最小值为 8
        width = max(width, 8)
        height = max(height, 8)
        
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    
    # 在开始处理之前记录开始时间
    start_time = datetime.datetime.now()  # 添加这一行以定义 start_time
    
    # 使用提供的参数生成图像
    with torch.no_grad():
        image = model(
            request.prompt,
            height=height,
            width=width,
            guidance_scale=3.5,
            num_inference_steps=request.num_inference_steps,
            max_sequence_length=512,
            generator=torch.Generator(device).manual_seed(0)
        ).images[0]
    
    # 生成唯一的文件名
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    random_number = random.randint(1000, 9999)
    image_name = f"{timestamp}_{random_number}.png"
    image_path = os.path.join("images", image_name)
    os.makedirs("images", exist_ok=True)
    image.save(image_path)
    
    # 计算文件大小
    file_size = os.path.getsize(image_path)  # 获取文件大小（字节）
    
    # 计算处理时间
    processing_time = (datetime.datetime.now() - start_time).total_seconds()  # 处理时间（秒）
    
    # 返回完整的 URL 和其他信息
    return {
        "image_url": f"{base_url}{image_name}",  # 使用配置文件中的 BASE_URL
        "file_size": file_size,  # 文件大小（字节）
        "processing_time": processing_time  # 处理时间（秒）
    }

@app.get("/images/{image_name}")
async def get_image(image_name: str):
    image_path = os.path.join("images", image_name)
    if not os.path.exists(image_path):
        raise HTTPException(status_code=404, detail="Image not found")
    return FileResponse(image_path)
